from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import yaml
import numpy as np
from dataclasses import is_dataclass, fields
from typing import Any, Dict, Type, TypeVar, get_type_hints

from ezgl import Matrix4x4


__all__ = [
    'SaveLoadConfig',
    'HandEyeConfig',
]


T = TypeVar('T')


def serialize_value(value: Any):
    """
    递归序列化单个值
    """
    if isinstance(value, np.ndarray):
        return {
            '_type_': 'np.ndarray',
            'dtype': str(value.dtype),
            'data': value.tolist() if value.size > 0 else [],
        }
    elif isinstance(value, Matrix4x4):
        return {
            '_type_': 'Matrix4x4',
            'data': value.toVector7d().tolist(),
        }
    return value

def deserialize_value(data) -> Any:
    """
    递归反序列化单个值
    """
    if isinstance(data, dict) and data.get('_type_') == 'np.ndarray':
        dtype = np.dtype(data['dtype'])
        return np.array(data['data'], dtype=dtype)
    
    elif isinstance(data, dict) and data.get('_type_') == 'Matrix4x4':
        return Matrix4x4.fromVector7d(*data['data'])

    return data

def serialize_dataclass(obj: Any) -> Dict[str, Any]:
    """
    将dataclass对象序列化为可存储为YAML的字典
    """
    if not is_dataclass(obj):
        raise TypeError(f"{type(obj)} is not dataclass")
    
    result = {}    
    for field in fields(obj):
        value = getattr(obj, field.name)

        if is_dataclass(value):
            result[field.name] = serialize_dataclass(value)
        else:
            result[field.name] = serialize_value(value)
            
    return result

def deserialize_dataclass(data: Dict[str, Any], cls: Type[T]) -> T:
    """
    从字典反序列化为dataclass对象
    """
    if not is_dataclass(cls):
        raise TypeError(f"{cls} is not dataclass")
    
    field_types = get_type_hints(cls)
    field_values = {}
    
    for field_name, field_value in data.items():
        if field_name not in field_types:
            print(f"Warning: Field '{field_name}' not found in {cls.__name__}, skipping.")
            continue  # 跳过未知字段
        
        if is_dataclass(field_types[field_name]):
            field_values[field_name] = deserialize_dataclass(field_value, field_types[field_name])

        else:
            field_values[field_name] = deserialize_value(field_value)

    return cls(**field_values)


class SaveLoadConfig:
    """
    Base configuration class.
    """
    def save(self, path: Union[str, Path]):
        """
        Save the configuration to a YAML file.
        """
        serialize_data = serialize_dataclass(self)
        if isinstance(path, str):
            path = Path(path)
            path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w', encoding='utf-8') as f:
            yaml.dump(serialize_data, f, default_flow_style=None, allow_unicode=True) 

    @classmethod
    def load(cls, path: Union[str, Path]):
        """
        load the yaml file to a dataclass.
        """
        with open(path, 'r', encoding='utf-8') as f:
            data = yaml.safe_load(f)
        if not isinstance(data, dict):
            raise ValueError("The yaml file is not a valid yaml file.")
        return deserialize_dataclass(data, cls)
        
    

#------------------------------------------------------------------------------

@dataclass
class HandEyeConfig(SaveLoadConfig):
    """
    Configuration for hand-eye calibration.
    """
    hand2eye: Matrix4x4 = field(default_factory=Matrix4x4)
    base2map: Matrix4x4 = field(default_factory=Matrix4x4)


if __name__ == "__main__":

    @dataclass
    class TestConfig(SaveLoadConfig):
        val: int = 1
        text: str = "hello"
        list_int: List[int] = field(default_factory=list)
        list_float: List[float] = field(default_factory=list)
        array: np.ndarray = field(default_factory=lambda: np.array([1, 2, 3], dtype=np.float32))
        nested: HandEyeConfig = HandEyeConfig()
    
    a = TestConfig(list_int=[1, 2], list_float=[1.0, 2.0])
    print(a)
    a.save("test.yaml")

    b = TestConfig.load("test.yaml")
    print(b)